package rock.registry.zk.util;

import lombok.extern.slf4j.Slf4j;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.imps.CuratorFrameworkState;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.framework.recipes.cache.PathChildrenCacheListener;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.zookeeper.CreateMode;
import rock.common.constants.RpcConstants;
import rock.common.enums.RpcConfigEnum;
import rock.common.util.PropertiesFileUtil;

import java.net.InetSocketAddress;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

@Slf4j
public class ZkCuratorUtil {

    private static CuratorFramework zkClient;
    private static final Map<String, List<String>> SERVICE_ADDRESS_MAP = new ConcurrentHashMap<>();
    private static final Set<String> REGISTERED_PATH_SET = ConcurrentHashMap.newKeySet();

    public ZkCuratorUtil() {
    }

    /**
     * 创建持久化节点，在client断连接后还会保存，temporary node断连接后消失
     * @param zkClient zk client
     * @param path node path
     */
    public static void createPersistNode(CuratorFramework zkClient, String path) {
        try {
            // check cache already exist or zk already exist
            if (REGISTERED_PATH_SET.contains(path) || zkClient.checkExists().forPath(path) != null) {
                log.info("node already exist, node is [{}]", path);
            } else {
                // not exist then create with CreateMode.PERSISTENT
                zkClient.create().creatingParentsIfNeeded().withMode(CreateMode.PERSISTENT).forPath(path);
                log.info("node was created success. node is [{}]", path);
            }
            // update cache always, keep newest
            REGISTERED_PATH_SET.add(path);
        } catch (Exception e) {
            log.error("create persist node for path [{}] fail.", path);
        }
    }

    /**
     * 获取所有服务 for serviceName
     * @param zkClient zk client
     * @param serviceName 服务名
     * @return 所有可用
     */
    public static List<String> getChildrenNodes(CuratorFramework zkClient, String serviceName) {
        if (SERVICE_ADDRESS_MAP.containsKey(serviceName)) {
            return SERVICE_ADDRESS_MAP.get(serviceName);
        }

        List<String> result = null;
        String servicePath = RpcConstants.ZK_REGISTER_ROOT_PATH + "/" + serviceName;
        try {
            // read from zk
            result = zkClient.getChildren().forPath(servicePath);
            // update cache map
            SERVICE_ADDRESS_MAP.put(serviceName, result);
            // register wather for serviceName
            registerWatcher(serviceName, zkClient);
        } catch (Exception e) {
            log.error("get children nodes for path [{}] fail.", servicePath);
        }
        return result;
    }

    /**
     * 清空注册中心数据
     * @param zkClient zk client
     * @param inetSocketAddress 服务地址
     */
    public static void clearRegistry(CuratorFramework zkClient, InetSocketAddress inetSocketAddress) {
        // 遍历删除zk path
        REGISTERED_PATH_SET.stream().parallel().forEach(p -> {
            try {
                if (p.endsWith(inetSocketAddress.toString())) {
                    zkClient.delete().forPath(p);
                }
            } catch (Exception e) {
                log.error("clear registry for path[{}] fail", p);
            }
        });
        log.info("all registered services on the server are cleared: [{}]", REGISTERED_PATH_SET.toString());
    }

    /**
     * 连接zk，并返回zkClient
     * @return
     */
    public static CuratorFramework getZkClient() {

        if (zkClient != null && zkClient.getState() == CuratorFrameworkState.STARTED) {
            return zkClient;
        }

        Properties properties = PropertiesFileUtil.readPropertiesFile(RpcConfigEnum.RPC_CONFIG_PATH.getPropertyVal());
        String zkAddress =
                properties != null && properties.getProperty(RpcConfigEnum.ZK_ADDRESS.getPropertyVal()) != null
                ? properties.getProperty(RpcConfigEnum.ZK_ADDRESS.getPropertyVal())
                : RpcConstants.DEFAULT_ZOOKEEPER_ADDRESS;

        RetryPolicy retryPolicy = new ExponentialBackoffRetry(RpcConstants.ZK_SLEEP_SEC, RpcConstants.ZK_MAX_RETRIES);
        zkClient = CuratorFrameworkFactory.builder()
                .connectString(zkAddress)
                .retryPolicy(retryPolicy)
                .build();
        zkClient.start();

        try {
            // wait timeout seconds until connect to zk, timeout throw exception
            if (!zkClient.blockUntilConnected(RpcConstants.WAIT_TIMEOUT_CONNECT_ZK_SEC, TimeUnit.SECONDS)) {
                throw new RuntimeException("Connect to ZK timeout.");
            }
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        return zkClient;
    }

    /**
     * 注册监听器 特定节点
     * @param serviceName 服务名称
     * @param zkClient zk
     * @throws Exception
     */
    private static void registerWatcher(String serviceName, CuratorFramework zkClient) throws Exception {
        String servicePath = RpcConstants.ZK_REGISTER_ROOT_PATH + "/" + serviceName;
        PathChildrenCache pathChildrenCache = new PathChildrenCache(zkClient, servicePath, true);
        PathChildrenCacheListener pathChildrenCacheListener = (curatorFramework, pathChildrenCacheEvent) -> {
            List<String > addresses = curatorFramework.getChildren().forPath(servicePath);
            SERVICE_ADDRESS_MAP.put(serviceName, addresses);
        };

        pathChildrenCache.getListenable().addListener(pathChildrenCacheListener);
        pathChildrenCache.start();
    }
}
